import torch
import numpy as np
from torch import nn

from np2c_header import np2c_header

c2 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=5, stride=1, padding=0)
din = torch.randn([1,1,32,32])

for name, p in c2.named_parameters():
    print(name, p.shape)
    np2c_header(p.detach().numpy(), name)

dout = c2(din)
print(dout.shape)

np2c_header(din.detach().numpy(), 'din')
np2c_header(dout.detach().numpy(), 'dout')